import os
import pika
import sys
import time
from pathlib import Path

# 配置参数
RABBITMQ_URL = 'amqp://user:password@172.16.168.58:5672/'  # RabbitMQ连接地址
QUEUE_NAME = 'imgurl'              # 队列名称
VALID_EXTENSIONS = ('.jpg', '.jpeg')  # 支持的图片扩展名

def scan_jpg_files(root_dir):
    """
    扫描目录获取所有JPG/JPEG文件名
    返回: (有效文件列表, 错误信息)
    """
    jpg_files = []
    
    try:
        # 验证目录有效性
        root_path = Path(root_dir)
        if not root_path.is_dir():
            return [], f"无效的目录: {root_dir}"

        # 遍历目录
        for entry in root_path.rglob('*'):
            # 跳过目录和隐藏文件
            if entry.is_dir() or entry.name.startswith('.'):
                continue
            
            # 检查扩展名（不区分大小写）
            if entry.suffix.lower() in VALID_EXTENSIONS:
                jpg_files.append(entry.name)
                
        return jpg_files, None
    
    except Exception as e:
        return [], f"扫描错误: {str(e)}"

def produce_messages(file_list):
    """
    发送消息到RabbitMQ队列
    """
    try:
        # 建立连接
        params = pika.URLParameters(RABBITMQ_URL)
        connection = pika.BlockingConnection(params)
        channel = connection.channel()
        
        # 声明持久化队列
        channel.queue_declare(
            queue=QUEUE_NAME,
            durable=True,    # 队列持久化
            exclusive=False,
            auto_delete=False
        )
        
        # 发送消息
        for filename in file_list:
            channel.basic_publish(
                exchange='',
                routing_key=QUEUE_NAME,
                body=filename,
                properties=pika.BasicProperties(
                    delivery_mode=2,  # 消息持久化
                    timestamp=int(time.time())
                )
            )
            print(f" [x] Sent {filename}")
            time.sleep(0.001)  # 1毫秒延迟
            
        connection.close()
        return True, None
        
    except Exception as e:
        return False, f"MQ错误: {str(e)}"

if __name__ == "__main__":
    # 检查命令行参数
    if len(sys.argv) < 2:
        print("请指定扫描目录，如: python jpgscan.py /path/to/images")
        sys.exit(1)
        
    target_dir = sys.argv[1]
    
    # 扫描文件
    files, error = scan_jpg_files(target_dir)
    if error:
        print(f"扫描失败: {error}")
        sys.exit(1)
        
    if not files:
        print("未找到任何JPG文件")
        sys.exit(0)
        
    print(f"找到 {len(files)} 个图片文件，开始发送...")
    
    # 发送消息
    success, error = produce_messages(files)
    if not success:
        print(f"消息发送失败: {error}")
        sys.exit(1)
        
    print("所有消息发送完成")

